package com.ktjiaoyu.travelaiagent.memory;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.StrUtil;
import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 基于Kryo序列化保存大模型对话记忆
 */
public class KryoChatMemory implements ChatMemory {
    /** 会话保存路径 */
    private String savePath;
    public KryoChatMemory(String savePath) {
        if(StrUtil.isBlank(savePath)) {
            throw new IllegalArgumentException("The save path cannot be null or empty.");
        }
        // 文件夹不存在，则创建文件夹，注意要用mkdirs方法
        if(!FileUtil.exist(savePath)) {
            new File(savePath).mkdirs();
        }
        this.savePath = savePath;
    }
    /** 创建Kryo实例方法 */
    private Kryo createKryo() {
        Kryo kryo = new Kryo();
        // 允许未注册的类被序列化，但可能导致性能下降
        kryo.setRegistrationRequired(false);
        // 该策略通过jvm api创建对象，这会创建一个完全空的对象（即不执行任何代码中的初始化工作）
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
        return kryo;
    }

    /** 反序列化会话消息列表 */
    private List<Message> readConversation(String conversationId, Kryo kryo) {
        File file = getConversationFile(conversationId);
        List<Message> messages = new ArrayList<>();
        if (FileUtil.exist(file)) {
            // 如果文件存在，则使用Kryo读取数据
            try (Input input = new Input(new FileInputStream(file))) {
                messages = kryo.readObject(input, ArrayList.class);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return messages;
    }

    /** 序列化会话消息列表 */
    private void writeConversation(String conversationId, List<Message> messages, Kryo kryo) {
        File file = getConversationFile(conversationId);
        try (Output output = new Output(new FileOutputStream(file))) {
            kryo.writeObject(output, messages);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /** 获取会话文件 */
    private File getConversationFile(String conversationId) {
        // 以conversationId为文件名，保存到savePath目录下
        return new File(savePath, conversationId + ".bin");
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        // Kryo并非线程安全，因此需要每次创建新的Kryo实例
        Kryo kryo = createKryo();
        // 先读取旧的会话消息
        List<Message> conversationMessages = readConversation(conversationId, kryo);
        // 添加新的消息
        conversationMessages.addAll(messages);
        // 重新写入会话消息
        writeConversation(conversationId, conversationMessages, kryo);
    }

    /** 获取最后N条会话消息 */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        Kryo kryo = createKryo();
        // 读取会话消息
        List<Message> conversationMessages = readConversation(conversationId, kryo);
        // 截取获取最后N条消息
        List<Message> messages = conversationMessages.subList(Math.max(0, conversationMessages.size() - lastN), conversationMessages.size());
        // 返回消息列表
        return messages;
    }

    @Override
    public void clear(String conversationId) {
        // 删除会话文件
        FileUtil.del(getConversationFile(conversationId));
    }
}
